# coding:utf-8


"""Config

Author:
    name: reeseimk
    email: reeseimk@163.com

Homepage: https://gitee.com/reeseimk/mindspore_bert
"""

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname("__file__"), "./")))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname("__file__"), "../")))
import ast
import argparse
from pprint import pformat
import yaml
import mindspore.common.dtype as mstype
from src.bert_model import BertConfig

class SetConfig(object):
    def __init__(self) -> None:
        self.currentDir = os.path.abspath(os.path.join(os.path.dirname("__file__"), "./"))
        self.parentDir = os.path.abspath(os.path.join(os.path.dirname("__file__"), "../"))
        
    def get_abs_path(self, path_relative: str=""):
        path = os.path.join(self.parentDir, path_relative)
        return path
    
    def read_yaml(self, config_path: str=""):
        with open(config_path, encoding="utf-8") as f:
            result = yaml.load_all(f.read(), Loader=yaml.FullLoader)  # loader可选择BaseLoader、SafeLoader、FullLoader、UnsafeLoader
            result = [x for x in result]
            return result
        
    def parse_cli_to_yaml(self, parser, cfg):
        """
        Parse command line arguments to the configuration according to the default yaml.

        Args:
            parser: Parent parser.
            cfg: Base configuration.
            helper: Helper description.
            cfg_path: Path to the default yaml config.
        """
        parser = argparse.ArgumentParser(description="[REPLACE THIS at config.py]", parents=[parser])
        for item in cfg:
            parser.add_argument("--" + item, type=type(cfg[item]), default=cfg[item])
        args = parser.parse_args()
        return args
    
    def get_config(self):
        config_path = self.get_abs_path(path_relative="config_dir/bert_config.yaml")
        yaml_res = self.read_yaml(config_path)
        # print(yaml_res)
        cfg_data = yaml_res[0]
        pretrain_load_ckpt_path = self.get_abs_path(path_relative=cfg_data['pretrain_load_ckpt_path'])
        vocab_path = self.get_abs_path(path_relative=cfg_data['vocab_file_path'])
        base_bert_cfg = cfg_data['base_bert_cfg']
        # print(pretrain_load_ckpt_path, vocab_path, base_bert_cfg)
        parser = argparse.ArgumentParser(description="default name", add_help=False)
        parser.add_argument("--pretrain_load_ckpt_path", default=pretrain_load_ckpt_path, help="pretrain_load_ckpt_path")
        parser.add_argument("--vocab_file_path", default=vocab_path, help="vocab_file_path")
        parser.add_argument("--base_bert_cfg", default=base_bert_cfg, help="vocab_path")
        args = self.parse_cli_to_yaml(parser, base_bert_cfg)
        bert_config = BertConfig(**base_bert_cfg)  # base_bert_cfg.type=dict
        bert_config.dtype = mstype.float32  # 特殊的，需要把str转为mstype.float32
        bert_config.compute_type = mstype.float16  # 特殊的，需要把str转为mstype.float16
        return args, bert_config
    

if __name__ == "__main__":
    set_config = SetConfig()
    # config_path = set_config.get_config_path(path_relative="config_dir/pretrain_config.yaml")
    # print(config_path)
    args, base_bert_cfg = set_config.get_config()
    print(args.vocab_file_path)
    print(base_bert_cfg, type(base_bert_cfg))
    print(base_bert_cfg.dtype, type(base_bert_cfg.dtype))